import pulp

import random, sys, operator, re, subprocess
from pylinpro import pylinpro
from matlib import matzero, matsetblock
from pprint import pprint
import progressbar
###############

glpk = pulp.GLPK(msg=0)

#===========================================================================#

def constructLP(candidates, winner, alpha=0.01, upBound=None, sumTo1=False,
        verbose=False, signDiffOnly=False, restriction=None):
    """
    Constructs the linear programming problem defined by a collection of
    violation vectors/rows and a given winner among them.  violnVecs should be
    a list of pairs of candidate names and tuples of integers (all of the same
    length). winner should be the winning (name,violnVec) candidate pair.
    alpha is a discretization constant, upBound is an optional upper bound to
    be imposed on constraint weights, and if sumTo1 is specified True, the
    additional constraint that all constraint weights sum to 1.0 will be
    imposed on the system. 

    Returns an instance of pulp.LpProblem, whose solve() method can be called
    to find a set of weights satisfying the system.
    """ 

    #print 'violnVecs=', violnVecs
    #print 'winner=', winner
    if verbose:
        print 'New LpProblem'
        print 'candidates:'
        pprint(candidates)
        print 'winner:'
        pprint(winner)

    prob = pulp.LpProblem("findWieghts", pulp.LpMinimize)

    # First we establish each weight as a LP variable
    vars = []
    winnerOutput, winnerViolns = winner
    for constrNum in range(len(winnerViolns)):
        # Each weight is constrained to be at least alpha
        # Each variable (weight) name is of the form 'wN'
        vars.append(pulp.LpVariable("w%s" % constrNum, lowBound=alpha,
            upBound=upBound))

    # sum iteratively applies +, which is overloaded for LpVariable
    prob += sum(vars)#, "obj" # our objective is to minimize sum of weights

    if sumTo1:
        prob += sum(vars) == 1.0

    if restriction:
        for coefs, intercept, ineq in restriction:
            row = []
            for c,w in enumerate(coefs):
                row.append(w*vars[c])
            if ineq == 0:
                prob += sum(row) == intercept
            elif ineq > 0:
                prob += sum(row) >= intercept
            elif ineq < 0:
                prob += sum(row) <= intercept

    for loserOutput, loserViolns in candidates:
        if (loserOutput, loserViolns) == (winnerOutput, winnerViolns):
            continue
        row = []
        for constrNum, winnerV in enumerate(winnerViolns):
            loserV = loserViolns[constrNum]
            diff = loserV - winnerV

            # Possibly truncate the difference to just its sign
            if signDiffOnly:
                if diff > 0.0:
                    diff = 1.0
                elif diff < 0.0:
                    diff = -1.0

            weight = vars[constrNum]
            row.append(diff*weight)

        if verbose:
            print '"%s" vs "%s" => %s' % (winnerOutput, loserOutput, sum(row) >= alpha)
        prob += sum(row) >= alpha#, loser

    return prob

#===========================================================================#

def addToLP(prob, candidates, winner, alpha=0.01, verbose=False,
        signDiffOnly=False):
    if verbose:
        print 'Adding to LpProblem'

    # This should order the problem variables by constraint index. We assume
    # that each variable name is of the form 'wN' where N is the string
    # representation of an integer -- so we order by the integers gotten by
    # lopping off the first char (the 'w').
    vars = sorted(prob.variables(), cmp = lambda x,y: cmp(int(str(x)[1:]),
        int(str(y)[1:])))
    winnerOutput, winnerViolns = winner

    for loserOutput, loserViolns in candidates:
        if (loserOutput, loserViolns) == (winnerOutput, winnerViolns):
            continue
        row = []
        for constrNum, winnerV in enumerate(winnerViolns):
            loserV = loserViolns[constrNum]
            diff = loserV - winnerV

            # Possibly truncate the difference to just its sign
            if signDiffOnly:
                if diff > 0.0:
                    diff = 1.0
                elif diff < 0.0:
                    diff = -1.0

            weight = vars[constrNum]
            row.append(diff*weight)

        if verbose:
            print '"%s" vs "%s" => %s' % (winnerOutput, loserOutput, sum(row) >= alpha)
        prob += sum(row) >= alpha#, loser

    return prob

#===========================================================================#

def constructComplexLP(tableaux, winners, alpha=0.01, upBound=None,
        sumTo1=False, verbose=False, signDiffOnly=False, restriction=None):

    if verbose:
        print 'constructComplexLP:\ntableaux:'
        pprint(tableaux)
        print 'winners:'
        pprint(winners)

    prob = constructLP(tableaux[0], winners[0], alpha, upBound, sumTo1,
            verbose, signDiffOnly=signDiffOnly, restriction=restriction)
    # TODO: upBound, sumTo1
    for candidates, winner in zip(tableaux[1:], winners[1:]):
        addToLP(prob, candidates, winner, alpha, verbose,
                signDiffOnly=signDiffOnly)
    return prob

#===========================================================================#

def randomTableau(candidates, numConstraints, minViol=0, maxViol=20):
    """
    Returns a random tableau (violnVecs dictionary) of the given candidates
    over the given number of constraints. Each candidate is given a random
    number of violations of each constraint, between minViol and maxViol, such
    that each candidate's violation vector is distinct, and each member of each
    pair of candidates is preferred by at least one constraint.

    """
    #FIXME: new tableau representation (list of (output,violns) tuples).
    violnVecs = {}
    for c in candidates:
        violnsReady = False
        violns = None
        while not violnsReady:
            violns = [random.randint(minViol, maxViol) for con in \
                xrange(numConstraints)]
            
            if all([(v==minViol) for v in violns]) or \
                    all([(v==maxViol) for v in violns]):
                # try again
                continue
            if not violnVecs:
                break
            else:
                violnsReady = True
                for otherC, otherViolns in violnVecs.iteritems():
                    gt = False
                    for con in xrange(numConstraints):
                        if violns[con] > otherViolns[con]:
                            gt = True
                            break
                    lt = False
                    for con in xrange(numConstraints):
                        if violns[con] < otherViolns[con]:
                            lt = True
                            break
                    if not (gt and lt):
                        violnsReady = False
                        break
        violnVecs[c] = violns

    return violnVecs

#===========================================================================#

def findWeightsLP(prob):
    """
    Tries to solve the given pulp problem instance with glpk; returns a
    dictionary from weight names to values if an optimal solution is found,
    otherwise returns False.
    """
    #print prob
    glpk.solve(prob)
    if pulp.LpStatus[prob.status] == 'Optimal':
        return dict([(v.name, v.varValue) for v in prob.variables()])
    elif pulp.LpStatus[prob.status] == 'Undefined':
        return False
    raise ValueError('Unexpected lp status!')

#===========================================================================#

def findWeights(candidates, winner, alpha=0.01, upBound=None, sumTo1=False):
    """
    Constructs the linear programming problem defined by a collection of
    violation vectors/rows and a given winner among them, and returns a
    dictionary describing a setting of constraint weights that renders the
    given winner optimal; returns False if no such setting is possible.

    Arguments are as for constructLP(). 
    """

    prob = constructLP(candidates, winner, alpha, upBound, sumTo1)
    return findWeights(prob)

#===========================================================================#

def isContenderIn(candidates, winner, alpha=0.01, upBound=None, sumTo1=False,
        signDiffOnly=False, oldWays=False):
    """
    Returns a boolean indicating whether there exists a setting of constraint
    weights that renders the indicated candidate a winner among the given
    collection of violation vectors/rows.

    Arguments are as for constructLP().
    """
    if oldWays:
        prob = constructLP(candidates, winner, alpha, upBound, sumTo1,
               signDiffOnly=signDiffOnly)
        glpk.solve(prob)
        if pulp.LpStatus[prob.status] == 'Optimal':
            return True
        elif pulp.LpStatus[prob.status] == 'Undefined':
            return False
        raise ValueError('Unexpected lp status!')

    rowIds = range(len(candidates))
    colIds = range(len(candidates[0][1]))
    winnerId = candidates.index(winner)
    diffMatrix = [tuple([b-a for a,b in zip(candidates[winnerId][1],
        candidates[r][1])]) for r in rowIds]

    #pprint(diffMatrix)

    while colIds:
        #print 'top of while. colIds =', colIds, 'rowIds =', rowIds
        removedCol = False
        colMins = [(c, min([diffMatrix[r][c] for r in rowIds])) for c in colIds]
        #print 'colMins =', colMins
        for c,m in colMins:
            if m >= 0:
                # this column has no Ls
                colIds.remove(c)
                removedCol = True

                # remove any row that contains a positive difference at column c
                rowIds = [r for r in rowIds if diffMatrix[r][c] <= 0]

        if not removedCol:
            subCandidates = [(candidates[r][0], tuple([candidates[r][1][c] for\
                c in colIds])) for r in rowIds]
            subWinner = (winner[0], tuple([winner[1][c] for c in colIds]))

            # check if the winner is pareto efficient (no need to solve linear
            # system if it's not)
            for cand, violns in subCandidates:
                if cand == subWinner[0]: continue
                if all([violns[i] - subWinner[1][i] <= 0 for i in
                        xrange(len(violns))]):
                    return False

            #print 'Sub problem: %r wins among %r' % (subWinner, subCandidates)

            if len(subCandidates) <= 2:
                # Only collective bounding differs from OT; this is a pairwise case.
                return False

            prob = constructLP(subCandidates, subWinner, alpha, upBound, sumTo1,
                    signDiffOnly=signDiffOnly)
            #prob = constructLP(candidates, winner, alpha, upBound, sumTo1,
            #        signDiffOnly=signDiffOnly)

            # THE GLPK WAY
            glpk.solve(prob)
            if pulp.LpStatus[prob.status] == 'Optimal':
                return True
            elif pulp.LpStatus[prob.status] == 'Undefined':
                return False
            raise ValueError('Unexpected lp status!')

            # The linpro way
            #return solvePyLinPro(prob)

    return True

#===========================================================================#

def contendersIn(violnVecs, alpha=0.01, upBound=None, sumTo1=False,
        signDiffOnly=False, oldWays=False):
    """
    Returns the set of members of the given collection of violation
    vectors/rows for which there exists a setting of constraint weights that
    renders that row optimal.

    violnVecs should be a list of candidate violation vectors (tuples of
    ints). Returns a set of violation vectors, without candidate names. This
    function is compatible with the interface expected by fsot.contenders for
    its contenderCostsFn argument.
    """
    violnVecs = list(set(violnVecs))
    if len(violnVecs) < 2:
        return violnVecs

    #print 'violnVecs:'
    #print violnVecs

    vecLen = max([len(v or ()) for v in violnVecs])
    candidates = [('row%s'%i, vec or (0,)*vecLen) for i,vec in enumerate(violnVecs)]
    return set([v for x,v in candidates if isContenderIn(candidates, (x,v),
        alpha=alpha, upBound=upBound, sumTo1=sumTo1, signDiffOnly=signDiffOnly,
        oldWays=oldWays)])

#===========================================================================#

#def contenders(eval, uRep, alpha=0.01, upBound=None, sumTo1=False,
#        signDiffOnly=False):
#    """
#    For the given eval (a composition of finite state constraints; see
#    fsot.makeEval()) and given input form uRep, returns a dictionary whose
#    keys are all of the non-harmonically-bounded candidate output forms for
#    that input, mapping to those candidates' violation vectors (tuples of
#    ints).
#    """
#    print 'contenders: %s' % uRep
#    ccf = lambda violnVecs: contendersIn(violnVecs, alpha, upBound, sumTo1,
#            signDiffOnly=signDiffOnly)
#    # TODO: weighting conditions
#    result = [(output, violns, None) for (output, violns, ercs) in \
#        fsot.contenders(eval, uRep, contenderCostsFn=ccf)]
#
#    # merge any tied contenders
#    resultNoTies = []
#    dead = []
#    for (o1, violns1, wc) in result:
#        if o1 in dead: continue
#        ties = []
#        for (o2, violns2, wc) in result:
#            if o2 in dead: continue
#            if o1 != o2 and violns2 == violns1:
#                ties.append(o2)
#                dead.append(o2)
#        if ties:
#            try:
#                resultNoTies.remove(('/'.join(ties), violns1, None))
#            except ValueError:
#                pass
#            resultNoTies.append(('/'.join(ties + [o1]), violns1, None))
#        else:
#            resultNoTies.append((o1, violns1, None))
#    return resultNoTies

#===========================================================================#

def concats(alphabet, k):
    """
    A simple recursive generator of all length-k concatenations of elements
    from the given alphabet.
    """
    if not k:
        yield ''
    else:
        alphabet = list(alphabet)
        for a in alphabet:
            for rest in concats(alphabet, k-1):
                yield a + rest

#===========================================================================#

def kleeneUnion(alphabet, k):
    for l in xrange(k+1):
        for n in concats(alphabet, l):
            yield n

#===========================================================================#

def findShatterableTableaux(numTableaux, numConstraints, attempts=100,
        showShatterable=False, showUnshatterable=False, showProblems=False):

    # FIXME: new tableau representation (list of (o,v)-tuples)
    numShattered = 0

    for attempt in xrange(1, attempts+1):
        print 'Attempt %d... ' % attempt,

        # Create binary tableaux, each with candidates named 'a' and 'b'
        tableaux = [randomTableau(['a', 'b'], numConstraints) \
                for t in xrange(numTableaux)]
        shattered = True
        # Enumerate all possible ways of setting winners across these tableaux
        for winners in concats(['a','b'], numTableaux):
            # For each winner set, construct the LP problem and see if it has
            # feasible weights
            prob = constructComplexLP(tableaux, winners)
            if findWeightsLP(prob):
                continue
            else:
                shattered = False
                print 'Unshatterable problem:'
                print prob
                break
        
        if shattered:
            print 'shatterable!'
            numShattered += 1
            if showShatterable:
                printTableaux(tableaux)
        else:
            print 'not shatterable.'
            if showUnshatterable:
                printTableaux(tableaux)

    print 'Summary: %d shatterable out of %d attempts.' % (numShattered,
            attempts)

#===========================================================================#

def printTableaux(tableaux, inputs=None, winners=None):
    inputs = inputs or ('?',)*len(tableaux)
    winners = winners or ('?',)*len(talbeaux)
    for tab, inp, win in zip(tableaux, inputs, winners):
        print '%s -> %s:' % (inp, win)
        for cand, violnVec in tab:
            print ('%s:\t' % cand) + '\t'.join([str(v) for v in violnVec])
        print

#===========================================================================#

def typologyWeighted(inputCandidates, verbose=False,
        returnWinnersChecked=False, alpha=0.01, progress=False,
        signDiffOnly=False, restriction=None):
    """
    inputCandidates should be a dictionary from input forms (strings) to
    contender (o,v,wc)-lists of the form returned by harmonic.contenders().
    """
    languages = {} # {Lg name: {input: output}}
    inputs = inputCandidates.keys()
    inputs.sort(cmp=lambda x,y: cmp(len(inputCandidates[x]),
        len(inputCandidates[y])), reverse=True)
    winnerSetsChecked = {} # {winner set: "Feasible" | "Infeasible"}

    firstInput = inputs[0]
    #weightVec = [0 for c in inputCandidates[firstInput][1][1]]
    #for coefs, intercept, ineq in restriction:
    #    if ineq == 0:
    #        for c, w in enumerate(coefs):
    #            if w == 1: 
    #                weightVec[c] = intercept
    #                break

    if progress:
        pbar = progressbar.ProgressBar(len(inputCandidates[firstInput])).start()
    else:
        pbar = None

    def visitWinner(input, output, violns, existingWinners):
        if verbose:
            print 'Visitng %r -> %r given %r' % (input, output, existingWinners)
        #if verbose:
        #    sys.stderr.write('.')

        # See if this input-output pair plus the existing winners is feasible
        orderedInputs = sorted(existingWinners.keys())
        tableaux = [[cand[:2] for cand in inputCandidates[input]]] + \
                [[cand[:2] for cand in inputCandidates[i]] for i in orderedInputs]
        winners = [(output, violns)] + [existingWinners[i] for i in orderedInputs]


        #print "\nTableaux:"
        #pprint(tableaux)
        #print "\nWinners:"
        #pprint(winners)
        prob = constructComplexLP(tableaux, winners, alpha=alpha,
                verbose=verbose, signDiffOnly=signDiffOnly,
                restriction=restriction)


        # if not, return, aborting this branch of the search
        if (not findWeightsLP(prob)):
            winnerSetsChecked[tuple(winners)] = 'Infeasible'
            #if verbose:
            #    sys.stderr.write('!')
            return

        winnerSetsChecked[tuple(winners)] = 'Feasible'

        # it was feasible... so now we add the feasible input-output pair to
        # the existing winners for this branch, pick another input, and visit
        # all the winners for that input
        existingWinners[input] = (output, violns)
        availInputs = [i for i in inputs if i not in\
                existingWinners]
        if not availInputs:
            # we've specified feasible winners for all inputs, which means we
            # have a language! add it to the languages dictionary and return
            # to the last branch
            #if verbose:
            #    sys.stderr.write('*')
            newLg = dict(existingWinners)
            for (oldLg, oldProb) in languages.values():
                if newLg == oldLg:
                    raise ValueError('Internal typology generation error: '\
                            'generated duplicate language.')
            lgName = 'L%04d' % (len(languages),)
            #print 'New language %s' % lgName
            #pprint(newLg)
            languages[lgName] = (newLg, prob) # copy this dict?
            return

        nextInput = availInputs[0]
        for (nextOutput, nextViolns, wc) in inputCandidates[nextInput]:
            visitWinner(nextInput, nextOutput, nextViolns, dict(existingWinners))

    # Do some recursive shit
    for (output, violns, wc) in inputCandidates[firstInput]:
        #print 'firstInput -> output: %s -> %s %s' % (firstInput, output,
        #        violns)
        visitWinner(firstInput, output, violns, {})
        if pbar:
            pbar.update(pbar.currval + 1)

    if pbar:
        pbar.finish()

    if returnWinnersChecked:
        return languages, winnerSetsChecked
    else:
        return languages

#===========================================================================#

def winnerUnderWeighting(tableau, weightVec):
    """
    S. T. A. N. K. Patrol Group
    """
    maxHarmony = None # None is minimal according to >
    winner = None
    for output, violnVec in tableau:
        harmony = sum(map(operator.mul, [-v for v in violnVec], weightVec))
        if harmony > maxHarmony:
            maxHarmony = harmony
            winner = (output, violnVec)
        elif harmony == maxHarmony:
            print 'Stank-ass tie!'
    return winner

#===========================================================================#

def tableauHasTiedWinners(tableau, weightVec):
    maxHarmony = None # None is minimal according to >
    for output, violnVec in tableau:
        harmony = sum(map(operator.mul, [-v for v in violnVec], weightVec))
        if harmony > maxHarmony:
            maxHarmony = harmony
        elif harmony == maxHarmony:
            return True
    return False

#===========================================================================#

def rankingToWeighting(ranking, maxViolations=99):
    """
    Returns a list w of weights, where w[i] the weighting of constraint i,
    such that the weights mimic an OT-style ranking of constraints; ranking is
    a list of constraint indices, in order of ranking (highest ranked to
    lowest).
    """
    weights = [None] * len(ranking)
    for i in range(-1, -len(ranking)-1, -1):
        constrIndex = ranking[i]
        if i == -1:
            w = 1.0
        else:
            w = (maxViolations+1)*weights[ranking[i+1]]
            for j in range(-1, i, -1):
                w += weights[ranking[j]]
        weights[constrIndex] = w
    return weights

#===========================================================================#

def solvePyLinPro(prob):
    """
    Given a PuLP linear programming problem object, solve it using pylinpro.py.
    Returns True iff the problem is feasible.
    """
    vars = sorted(prob.variables())
    N = len(vars)
    M = len(prob.constraints)
    JLIM = N + M + 1
    Basis = range(JLIM)
    AA = matzero(M + 2, JLIM)

    rows = []
    for cName, con in prob.constraints.iteritems():
        varCoeffs = dict([(v,0.0) for v in vars])
        varCoeffs.update(con)
        rows.append(0.0)
        for v in vars:
            rows.append(varCoeffs[v])
    rows += [0.0] + [1.0 for v in vars]
    matsetblock(AA, 0, 0, M, N, rows)

    LPcode = pylinpro(AA, M, N, 0, 0, M, False, 0, Basis, 0)

    if LPcode == 0:
        return True
    elif LPcode == 4:
        return False
    raise ValueError('Crazy LPcode %s!' % LPcode)

#===========================================================================#

def lpToPoly(prob, polyfile, lbound=0.0, ubound=1.0):
    """
    Converts an PuLP linear programming problem representation of an HG model
    (set of tableaux with winners specified) into a description of the
    weighting polytope consistent with the model, in polymake format. polyfile
    is the file-like object to write the description to.
    """
    f = polyfile
    f.write('INEQUALITIES\n')
    vars = sorted(prob.variables())
    # constrain each var to more than lbound
    for i in xrange(len(vars)):
        f.write('%s ' % lbound)
        for j in xrange(len(vars)):
            f.write(i==j and '1 ' or '0 ')
        f.write('\n')
    # constrain each var to less than ubound
    for i in xrange(len(vars)):
        f.write('%s ' % ubound)
        for j in xrange(len(vars)):
            f.write(i==j and '-1 ' or '0 ')
        f.write('\n')
    # inequalities from lp constraints
    for cName, con in prob.constraints.iteritems():
        varCoeffs = dict([(v,0) for v in vars])
        varCoeffs.update(con)
        f.write('0 ')
        for v in vars:
            f.write('%s ' % varCoeffs[v])
        f.write('\n')

#===========================================================================#

volumePat = re.compile(r'.*VOLUME\n(\d+)/(\d+)', re.DOTALL)

def problemVolume(prob, lowBound=0.0, upBound=1.0):
    f = open('weightVolume.poly', 'wb')
    lpToPoly(prob, f, lowBound, upBound)
    f.close()
    polymakeOut = open('polymake.out', 'wb')
    polymakeErr = open('polymake.err', 'wb')
    subprocess.call('polymake weightVolume.poly VOLUME', shell=True,
            stdout=polymakeOut, stderr=polymakeErr)
    polymakeOut.close()
    polymakeErr.close()
    output = open('polymake.out', 'r').read()
    match = volumePat.match(output)
    if not match:
        volume = 0.0 # bad polytope?
    else:
        numerator, denominator = match.groups()
        volume = float(numerator) / float(denominator)
    return volume

def weightVolume(tableaux, winners, lowBound=0.0, upBound=1.0, sumTo1=False,
        verbose=False, alpha=0.01):
    for winner, tab in zip(winners, tableaux):
        if winner not in tab:
            raise ValueError('Winner %s is not in its tableau %s.' % (winner,
                tableaux))
    prob = constructComplexLP(tableaux, winners, sumTo1=sumTo1, verbose=verbose,
            alpha=alpha)
    #if not findWeightsLP(prob):
    #    raise ValueError('Got infeasible problem.')
    #    return 0.0

    volume = problemVolume(prob, lowBound, upBound)

    if not volume:
        print 'Warning: 0 volume, infeasible.'
        print 'Winners:'
        print winners
        print 'Tableaux:'
        print tableaux
    return volume

def typologyWeightVolumes(typology, contenderDict):
    langVols = {}
    for language in typology:
        inputs = sorted(typology[language].keys())
        tableaux = []
        winners = []
        for input in inputs:
            tableau = []
            for o,v,e in contenderDict[input]:
                tableau.append((o,v))
            winner = typology[language][input]
            tableaux.append(tableau)
            winners.append(winner)

        #print 'Language %s:' % language
        #printTableaux(tableaux, inputs, winners)
        #break

        vol = weightVolume(tableaux, winners)
        langVols[language] = vol

    return langVols
